

segrap_task002_one_hot_label_names = {
    "GTVp": 1,
    "GTVnd": 2}

from light_training.preprocessing.preprocessors.preprocessor_multiinput_and_region_01norm_first import MultiInputAndRegionPreprocessor 
import numpy as np 
import pickle 
import json 


def process_train():
    base_dir = "./data/raw_data/"
    image_dir = "SegRap2023_Training_Set_120cases_task2"
    data_filenames = ["image.nii.gz", "image_contrast.nii.gz"]

    seg_filename = "seg_task2.nii.gz"
    preprocessor = MultiInputAndRegionPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    data_filenames=data_filenames,
                                    seg_filename=seg_filename,
                                    norm_clip_min=-175,
                                    norm_clip_max=250
                                   )
    
    out_spacing = [3.0, 0.54199219, 0.54199219]
    # out_spacing = [3.0, 1.0, 1.0]
    output_dir = "./data/fullres_task2/train/"

    preprocessor.run(output_spacing=out_spacing, 
                     output_dir=output_dir, 
                     all_labels_dict=segrap_task002_one_hot_label_names,
                     num_processes=32,
                    )

def process_val():
    # fullres spacing is [0.5        0.70410156 0.70410156]
    # median_shape is [602.5 516.5 516.5]
    base_dir = "./data/raw_data/Val"
    image_dir = "img"
    preprocessor = DefaultPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    label_dir=None,
                                   )

    out_spacing = [0.5, 0.70410156, 0.70410156]

    with open("./data_analysis_result.txt", "r") as f:
        content = f.read().strip("\n")
        print(content)
    content = eval(content)
    foreground_intensity_properties_per_channel = content["intensity_statistics_per_channel"]

    output_dir = "./data/fullres/val_test/"
    preprocessor.run(output_spacing=out_spacing, 
                     output_dir=output_dir,
                     all_labels=[1, ],
                     foreground_intensity_properties_per_channel=foreground_intensity_properties_per_channel,
                     num_processes=16)

def process_val_semi():
    # fullres spacing is [0.5        0.70410156 0.70410156]
    # median_shape is [602.5 516.5 516.5]
    base_dir = "./data/raw_data/Val_semi_postprocess"
    image_dir = "img"
    preprocessor = DefaultPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    label_dir="gt",
                                   )

    out_spacing = [0.5, 0.70410156, 0.70410156]

    with open("./data_analysis_result.txt", "r") as f:
        content = f.read().strip("\n")
        print(content)
    content = eval(content)
    foreground_intensity_properties_per_channel = content["intensity_statistics_per_channel"]

    output_dir = "./data/fullres/val_semi_postprocess/"
    preprocessor.run(output_spacing=out_spacing, 
                     output_dir=output_dir,
                     all_labels=[1, ],
                     foreground_intensity_properties_per_channel=foreground_intensity_properties_per_channel)


def plan():
    base_dir = "./data/raw_data/"
    image_dir = "SegRap2023_Training_Set_120cases_task2"
    data_filenames = ["image.nii.gz", "image_contrast.nii.gz"]

    seg_filename = "seg_task2.nii.gz"
    preprocessor = MultiInputAndRegionPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    data_filenames=data_filenames,
                                    seg_filename=seg_filename,
                                    norm_clip_min=-175,
                                    norm_clip_max=250
                                   )

    preprocessor.run_plan()

if __name__ == "__main__":

    # plan()

    process_train()
    # import time 
    # s = time.time()
    # process_val()
    # e = time.time()

    # print(f"preprocessing time is {e - s}")

 